// Copyright (c) 2020-present, INSPUR Co, Ltd. All rights reserved.
// This source code is licensed under Apache 2.0 License.

#include <chrono>
#include <iostream>
#include <string>
#include "pure_mem/key_index/abstract_tree.h"
#include "pure_mem/key_index/art/rowex_tree.h"
#include "pure_mem/key_index/art/rowex_tree_fullkey.h"
#include "util/testharness.h"

namespace rocksdb {

Slice genKey(int num) {
  std::string aa = std::to_string(num);
  char *ret = new char[aa.length() + 1];
  memcpy(ret, aa.c_str(), aa.length());
  ret[aa.length()] = 0;
  return Slice(ret, aa.length());
}

uint64_t loadKey(const void *tid) {
  std::string key((char *)tid);
  return stoi(key);
}

void loadKeyFromVoid(void *node, Slice &key) {
  std::string kkey((char *)node);
  key.data_ = (const char *)node;
  key.size_ = kkey.length();
}

void printTID(void *tid) {
  if (tid == nullptr) {
    std::cout << "printTID--------current tid is empty." << std::endl;
  } else {
    std::cout << "printTID--------key: " << loadKey(tid) << std::endl;
  }
}

class TreeExample {
public:
  ITree *tree_;

public:
  TreeExample(ITree *tree) : tree_(tree) {}
  void assertWithKey(void *t, int num) {
    if (t == nullptr) {
      //      std::cout << "error: t is empty for key " << num << std::endl;
      // exit(0);
      assert(num == 0);
      return;
    }

    int tKey = loadKey(t);
    //    std::cout << "assertWithKey  a: " << num << ", t: " << tKey <<
    //    std::endl;

    assert(tKey == num);
  }

  void insertWithNum(int num) {
    Slice curkey = genKey(num);
    void *next = nullptr;
    bool insert = tree_->insertNoReplace(curkey, (void *)(curkey.data_), next);
    assert(insert);
  }

  void *searchWithNum(int num) {
    Slice a = genKey(num);
    int load = loadKey(a.data_);
    void *val = tree_->getGE(a);
    // std::cout << "search for key: " << num << ", with result tid: " << val
    //           << std::endl;

    return val;
  }

  void testInsertAndGetLG() {
    void *next;
    Slice a5 = genKey(5);
    int loada5 = loadKey(a5.data_);
    bool insertSuc = tree_->insertNoReplace(a5, (void *)(a5.data_), next);
    // std::cout<< "----------a5: "<< a5 << ", cur: " <<cur << ", next: " <<
    // next <<std::endl;
    assert(insertSuc);
    assert(next == nullptr);

    Slice a5_1 = genKey(5);
    insertSuc = tree_->insertNoReplace(a5_1, (void *)(a5_1.data_), next);
    assert(!insertSuc);
    assert(next == a5.data_);

    Slice a6 = genKey(6);
    insertSuc = tree_->insertNoReplace(a6, (void *)(a6.data_), next);
    assert(insertSuc);
    assert(next == nullptr || next == a5.data_);

    Slice a2 = genKey(2);
    insertSuc = tree_->insertNoReplace(a2, (void *)(a2.data_), next);
    // std::cout<< "----------a2: "<< a2 << ", cur: " <<cur << ", next: " <<
    // next <<std::endl;
    assert(insertSuc);
    assert(next == a5.data_);

    Slice a4 = genKey(4);
    insertSuc = tree_->insertNoReplace(a4, (void *)(a4.data_), next);
    // std::cout<< "----------a4: "<< a4 << ", cur: " <<cur << ", next: " <<
    // next <<std::endl;
    assert(insertSuc);
    assert(next == a5.data_);

    void *b = 0;
    for (int i = 0; i < 100; i++) {
      Slice a = genKey(799 - i);
      insertSuc = tree_->insertNoReplace(a, (void *)(a.data_), next);
      //   std::cout << "----------a: " << a.ToString() << ", next: " << next
      //             << ",b: " << b << std::endl;
      assert(insertSuc);
      assert(next == b || next == a6.data_);
      b = (void *)(a.data_);
    }
  }

  void testNullorEmpty() {
    std::cout << "test case: key is null or empty!" << std::endl;

    printTID(searchWithNum(5));
    printTID(searchWithNum(-1));
    void *next;

    Slice a55 = genKey(55);
    bool insertSuc = tree_->insertNoReplace(a55, (void *)(a55.data_), next);
    // std::cout << "----------a55: " << a55.ToString() << ", next: " << next
    // << std::endl;
    assert(insertSuc);
    assert(next == nullptr);

    printTID(searchWithNum(5));
    printTID(searchWithNum(-1));

    Slice a_1 = genKey(-1);
    insertSuc = tree_->insertNoReplace(a_1, (void *)(a_1.data_), next);
    // std::cout << "----------a_1: " << a_1.ToString() << ", next: " << next
    //           << std::endl;
    assert(insertSuc);
    assert(next == a55.data_);

    printTID(searchWithNum(5));
    printTID(searchWithNum(-1));

    Slice a5 = genKey(5);
    insertSuc = tree_->insertNoReplace(a5, (void *)(a5.data_), next);
    // std::cout << "----------a5: " << a5.ToString() << ", next: " << next
    //           << std::endl;
    assert(insertSuc);
    assert(next == a55.data_);
  }

  void testEG() {
    std::cout << "test case: equal or greater!" << std::endl;

    insertWithNum(-1);
    assertWithKey(searchWithNum(-1), -1);
    assert(searchWithNum(1) == nullptr);
    insertWithNum(1);
    assertWithKey(searchWithNum(1), 1);
    insertWithNum(11);
    assertWithKey(searchWithNum(1), 1);
    assertWithKey(searchWithNum(11), 11);
    for (int i = 12; i <= 70; i++) {
      insertWithNum(i);
    }
    assertWithKey(searchWithNum(-1), -1);
    assertWithKey(searchWithNum(2), 20);
    insertWithNum(2);
    assertWithKey(searchWithNum(2), 2);
    assert(searchWithNum(9) == nullptr);
    assertWithKey(searchWithNum(10), 11);
    assertWithKey(searchWithNum(3), 30);
    assertWithKey(searchWithNum(44), 44);
  }

  void testPathCompress() {
    std::cout << "test case: path compress!" << std::endl;
    insertWithNum(1234);
    assertWithKey(searchWithNum(12), 1234);
    insertWithNum(12);
    assertWithKey(searchWithNum(1), 12);
    assertWithKey(searchWithNum(123), 1234);
    insertWithNum(11);
    assertWithKey(searchWithNum(122), 1234);
    assertWithKey(searchWithNum(111), 12);
    assertWithKey(searchWithNum(13), 0);
  }

  void testMultiPathCompress() {
    std::cout << "test case: path compress!" << std::endl;
    insertWithNum(1);
    assertWithKey(searchWithNum(1), 1);
    insertWithNum(10);
    assertWithKey(searchWithNum(1), 1);
    assertWithKey(searchWithNum(10), 10);
    insertWithNum(100);
    assertWithKey(searchWithNum(1), 1);
    assertWithKey(searchWithNum(10), 10);
    assertWithKey(searchWithNum(100), 100);
    insertWithNum(1000);
    assertWithKey(searchWithNum(1), 1);
    assertWithKey(searchWithNum(10), 10);
    assertWithKey(searchWithNum(100), 100);
    assertWithKey(searchWithNum(1000), 1000);

    insertWithNum(9000);
    assertWithKey(searchWithNum(9000), 9000);
    insertWithNum(900);
    assertWithKey(searchWithNum(900), 900);
    assertWithKey(searchWithNum(9000), 9000);
    insertWithNum(90);
    assertWithKey(searchWithNum(9000), 9000);
    assertWithKey(searchWithNum(90), 90);
    assertWithKey(searchWithNum(900), 900);
    insertWithNum(9);
    assertWithKey(searchWithNum(9), 9);
    assertWithKey(searchWithNum(90), 90);
    assertWithKey(searchWithNum(900), 900);
    assertWithKey(searchWithNum(9000), 9000);
  }

  static void get(TreeExample *te, uint64_t sum, int64_t begin) {
    {
      // Lookup
      auto starttime = std::chrono::system_clock::now();

      for (uint64_t i = 0; i < sum; i++) {
        auto val = te->tree_->getGE((char *)(genKey(i + begin).data_));
        if (loadKey(val) != (i + begin)) {
          std::cout << "wrong key read: " << val << " expected:" << (i + begin)
                    << std::endl;
          throw;
        }
      }
      auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
          std::chrono::system_clock::now() - starttime);
      printf("lookup,%ld,%f\n", sum, (sum * 1.0) / duration.count());
    }
  }

  static void insert(TreeExample *te, uint64_t sum, int64_t begin,
                     int treeType) {
    // uint64_t n = std::atoll(argv[1]);
    void **values = new void *[sum];

    // Generate values
    for (uint64_t i = 0; i < sum; i++)
      values[i] = (void *)(genKey(i + begin).data_);
    if (treeType > 0)
      std::random_shuffle(values, values + sum);

    printf("operation,n,ops/s\n");

    // Build tree_
    {
      auto starttime = std::chrono::system_clock::now();

      void *next;
      uint64_t succNum = 0;
      for (uint64_t i = 0; i < sum; i++) {

        bool insertSuc =
            te->tree_->insertNoReplace((char *)values[i], values[i], next);
        if (insertSuc)
          succNum++;
      }

      auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
          std::chrono::system_clock::now() - starttime);
      printf("insert,%ld,%f,success insert:%ld\n", sum,
             (sum * 1.0) / duration.count(), succNum);
    }
  }

  void multithreaded(uint64_t n, int treeType) {
    std::cout << "multi threaded:" << std::endl;

    // uint64_t n = std::atoll(argv[1]);
    void **values = new void *[n];

    // Generate values
    for (uint64_t i = 0; i < n; i++)
      values[i] = (void *)(genKey(i).data_);
    if (treeType > 0)
      // dense, random
      std::random_shuffle(values, values + n);

    printf("operation,n,ops/s\n");

    // Build tree_
    {
      auto starttime = std::chrono::system_clock::now();

      void *next;
      for (uint64_t i = 0; i < n; i++) {

        bool insertSuc =
            tree_->insertNoReplace((char *)values[i], values[i], next);
        assert(insertSuc);
        void *val = tree_->getGE((char *)values[i]);
        if (val != values[i]) {
          std::cout << "wrong key read: " << val << " expected:" << values[i]
                    << std::endl;
          throw;
        }
      }

      auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
          std::chrono::system_clock::now() - starttime);
      printf("insert,%ld,%f\n", n, (n * 1.0) / duration.count());
    }

    {
      // Lookup
      auto starttime = std::chrono::system_clock::now();

      for (uint64_t i = 0; i < n; i++) {
        auto val = tree_->getGE((char *)values[i]);
        if (val != values[i]) {
          std::cout << "wrong key read: " << val << " expected:" << values[i]
                    << std::endl;
          throw;
        }
      }
      auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
          std::chrono::system_clock::now() - starttime);
      printf("lookup,%ld,%f\n", n, (n * 1.0) / duration.count());
    }

    delete[] values;
  }
};

class TreeTest : public testing::Test, public testing::WithParamInterface<int> {
public:
  ITree *genTree(int n) {
    switch (n) {
    case 1:
      return new RowexTreeFullKey(loadKeyFromVoid, nullptr);
      break;
    case 2:
      return new RowexTree(loadKeyFromVoid);
      break;

    default:
      break;
    }
    return nullptr;
  }

  void deleteTree(int n, ITree *tree) {
    switch (n) {
    case 1:
      delete dynamic_cast<RowexTreeFullKey *>(tree);

      break;
    case 2:
      delete dynamic_cast<RowexTree *>(tree);
      break;

    default:
      break;
    }
  }
};

INSTANTIATE_TEST_CASE_P(artTreType, TreeTest, testing::Values(1, 2));

TEST_P(TreeTest, testNullorEmpty) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.testNullorEmpty();
  deleteTree(treeType, tree);
}
TEST_P(TreeTest, testInsertAndGetLG) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.testInsertAndGetLG();
  deleteTree(treeType, tree);
}
TEST_P(TreeTest, testEG) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.testEG();
  deleteTree(treeType, tree);
}
TEST_P(TreeTest, testPathCompress) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.testPathCompress();
  deleteTree(treeType, tree);
}
TEST_P(TreeTest, testMultiPathCompress) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.testMultiPathCompress();
  deleteTree(treeType, tree);
}

TEST_P(TreeTest, singlethreadedOrdered) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.multithreaded(1000000, 0);
  deleteTree(treeType, tree);
}

TEST_P(TreeTest, singlethreadedUnOrdered) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  rt.multithreaded(1000000, 1);
  deleteTree(treeType, tree);
}

TEST_P(TreeTest, multithreadedUnOrdered) {
  int treeType = GetParam();
  ITree *tree = genTree(treeType);
  TreeExample rt(tree);
  std::thread test1(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test2(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test3(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test4(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test5(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test6(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test7(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test8(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test9(TreeExample::insert, &rt, 1000000, 0, 1);
  std::thread test10(TreeExample::insert, &rt, 1000000, 0, 1);

  test1.join();
  test2.join();
  test3.join();
  test4.join();
  test5.join();
  test6.join();
  test7.join();
  test8.join();
  test9.join();
  test10.join();
  TreeExample::get(&rt, 1000000, 0);
  deleteTree(treeType, tree);
}
} // namespace rocksdb

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
  system("pause");
}
